from argparse import ArgumentParser

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import os.path as osp


def contrast_stretch(r, m, E, eps=1e-12):
    """Apply contrast stretching given by formula `s = T(r) = 1 / (1 + (m/r) ^ E)`,
    where `r` is the input ndarray, `s` is the output ndarray, `m` and `E` are two
    parameters that control the shape of the contrast stretching curve.
    """

    r = np.where(np.array(r) == 0.0, eps, np.array(r))

    return 1.0 / (1.0 + (m / r) ** E)


def unit_test_problem_01(m=0.5, E=5.0):
    """Problem 01: calculate the output of function `contrast_stretch` over 1000 sampled
    points in range [0, 1], and plot the contrast stretching curve."""

    print(f"Problem 01: plot contrast stretching curve over 1000 sampled points.")

    X = np.linspace(0.0, 1.0, 1000)
    Y = contrast_stretch(X, m, E)

    fig, ax = plt.subplots(figsize=(8.0, 8.0), num=1)
    fig.subplots_adjust(0.16, 0.14, 0.90, 0.88)

    ax.plot(X, Y, linewidth=1, label=f"m={m:.2f}, E={E:.2f}")
    ax.grid(visible=True, which="both", axis="both", linestyle="--")
    ax.set_title(f"s = T(r) = 1 / (1 + ({m}/r) ^ {E})", fontsize=14)
    ax.set_xlabel("r", fontsize=12)
    ax.set_ylabel("s", fontsize=12)

    fig_path = osp.join(osp.dirname(__file__), f"contrast_stretch-m_{m}_E_{E}.png")
    plt.savefig(fig_path)

    print(f'[ m = {m}, E = {E} ] => "{fig_path}"')


def unit_test_problem_02(e_min=1.0, e_max=9.0, num=10, sampler=None):
    """Problem 02: with fixed `m = 0.5`, change the value of `E`, plot the curves."""

    print(
        f"Problem 02: with fixed `m = 0.5`, change the value of `E`, plot the curves."
    )

    X = np.linspace(0.0, 1.0, 1000)
    E = np.linspace(e_min, e_max, num)
    if sampler:
        E = sampler(E)

    fig, ax = plt.subplots(figsize=(8.0, 8.0), num=2)
    fig.subplots_adjust(0.16, 0.14, 0.90, 0.88)

    colors = mpl.colormaps["rainbow"]

    for edx, e in enumerate(E):
        Y = contrast_stretch(X, m=0.5, E=e)
        ax.plot(X, Y, linewidth=1, label=f"E={e:.2f}", color=colors(edx / len(E)))
    ax.grid(visible=True, which="both", axis="both", linestyle="--")
    ax.set_title(f"s = T(r) = 1 / (1 + (0.5/r) ^ E)", fontsize=14)
    ax.set_xlabel("r", fontsize=12)
    ax.set_ylabel("s", fontsize=12)
    ax.legend(fontsize=8, ncols=2)

    fig_path = osp.join(osp.dirname(__file__), f"contrast_stretch-m_0.5.png")
    plt.savefig(fig_path)

    print(f'[ m = 0.5, varying E ] => "{fig_path}"')


curve_history = []  # Preserve curve parameters


def plot_curve(m_value, E_value, add_plot):
    """Plot contrast stretching curve in a matplotlib Figure."""

    X = np.linspace(0.0, 1.0, 1000)

    if not add_plot:
        curve_history.clear()
    curve_history.append([m_value, E_value, np.random.random()])

    fig, ax = plt.subplots(figsize=(8.0, 8.0), num=3, clear=True)
    fig.subplots_adjust(0.16, 0.14, 0.90, 0.88)

    colors = mpl.colormaps["rainbow"]

    for m, E, c_seed in curve_history:
        Y = contrast_stretch(X, m=m, E=E)
        ax.plot(
            X,
            Y,
            linewidth=1,
            label=f"m={m:.2f}, E={E:.2f}",
            color=colors(c_seed),
        )

    ax.grid(visible=True, which="both", axis="both", linestyle="--")
    ax.set_title(f"s = T(r) = 1 / (1 + (m/r) ^ E)", fontsize=14)
    ax.set_xlabel("r", fontsize=12)
    ax.set_ylabel("s", fontsize=12)
    ax.legend(fontsize=8, ncols=1)

    return fig


def generate_playback():
    """Generate a video using the history of the drawn contrast stretching curves."""
    import matplotlib.animation as man

    X = np.linspace(0.0, 1.0, 1000)

    fig, ax = plt.subplots(figsize=(8.0, 8.0), num=3, clear=True)
    fig.subplots_adjust(0.16, 0.14, 0.90, 0.88)

    colors = mpl.colormaps["rainbow"]

    artists, cum_artists = [], []

    for m, E, c_seed in curve_history:
        Y = contrast_stretch(X, m=m, E=E)
        lines = ax.plot(
            X,
            Y,
            linewidth=1,
            label=f"m={m:.2f}, E={E:.2f}",
            color=colors(c_seed),
            animated=True,
        )
        cum_artists.extend(lines)
        artists.append(cum_artists.copy())

    ax.grid(visible=True, which="both", axis="both", linestyle="--")
    ax.set_title(f"s = T(r) = 1 / (1 + (m/r) ^ E)", fontsize=14)
    ax.set_xlabel("r", fontsize=12)
    ax.set_ylabel("s", fontsize=12)
    ax.legend(fontsize=8, ncols=1)

    ani = man.ArtistAnimation(fig, artists, interval=500, blit=True)
    ani_path = osp.join(osp.dirname(__file__), f"contrast_stretch_replay.mp4")
    ani.save(ani_path)

    return ani_path


def start_graphical_demo():
    """Start GUI demo for parameter exploration and playback."""

    import gradio as gr

    with gr.Blocks(title="Contrast Stretching Demo") as demo:
        gr.Markdown(
            r"""
            ## 对比度拉伸变换图形化交互Demo

            对比度拉伸变换 $s = T(r) = \frac{1}{1 + (\frac{m}{r}) ^ E}$，
            其中 $s, r$ 分别表示输出与输入图像的灰度级，而 $m, E$ 是2个控制参数。
            """
        )

        with gr.Tab("Interactive"):
            with gr.Row():
                m_value = gr.Number(0.5, label="m", show_label=True)
                E_value = gr.Number(5.0, label="E", show_label=True)
            add_plot = gr.Checkbox(
                False, label="Add a New Curve to the Current Plot.", show_label=True
            )

            plot = gr.Plot(label="Contrast Stretching Curve", show_label=True)
            button_1 = gr.Button("Plot")
            button_1.click(
                plot_curve, inputs=[m_value, E_value, add_plot], outputs=plot
            )

        with gr.Tab("Playback"):
            video_preview = gr.Video(
                label="History Playback", show_label=True, format="mp4"
            )
            button_2 = gr.Button("Replay")
            button_2.click(generate_playback, outputs=video_preview)

    demo.launch()


def main_procedure(cmdargs):
    if not cmdargs.graphics:
        unit_test_problem_01(m=0.5, E=5.0)
        unit_test_problem_02(1.0, 10.0, 10)
    else:
        print("Start GUI demo [using backend: gradio] ...")
        start_graphical_demo()


if __name__ == "__main__":
    parser = ArgumentParser(description="Contrast Stretching Curve Demo.")

    parser.add_argument(
        "-g",
        "--graphics",
        default=False,
        action="store_true",
        help="Use graphical UI for interactive exploration.",
    )

    main_procedure(parser.parse_args())
